package job_types

import (
	"context"
	"fmt"
	"time"

	"gitlab.ctyuncdn.cn/tai/infra/cloud-controller/internal/job_manager"
	"gitlab.ctyuncdn.cn/tai/infra/cloud-controller/pkg/mysql"
)

// NotebookJobTypeHandler handles notebook jobs
type NotebookJobTypeHandler struct {
	*BaseJobTypeHandler
}

// NewNotebookJobTypeHandler creates a new notebook job type handler
func NewNotebookJobTypeHandler() *NotebookJobTypeHandler {
	base := NewBaseJobTypeHandler(job_manager.JobTypeNotebook, "Jupyter Notebook job handler")

	// Add notebook-specific parameters
	notebookParams := []job_manager.Parameter{
		{
			Name:        "notebook_path",
			Type:        "string",
			Required:    true,
			Description: "Path to the notebook file",
			Validation:  "notebook_path != ''",
		},
		{
			Name:        "kernel",
			Type:        "string",
			Required:    false,
			Default:     "python3",
			Description: "Jupyter kernel to use",
			Validation:  "kernel in ['python3', 'python2', 'r', 'julia']",
		},
		{
			Name:        "port",
			Type:        "int",
			Required:    false,
			Default:     8888,
			Description: "Port for Jupyter notebook server",
			Validation:  "port >= 1024 && port <= 65535",
		},
		{
			Name:        "token",
			Type:        "string",
			Required:    false,
			Default:     "",
			Description: "Jupyter notebook token for authentication",
		},
		{
			Name:        "allow_origin",
			Type:        "string",
			Required:    false,
			Default:     "*",
			Description: "Allowed origins for CORS",
		},
	}

	// Combine base and notebook-specific parameters
	allParams := append(base.supportedParameters, notebookParams...)
	base.supportedParameters = allParams

	return &NotebookJobTypeHandler{
		BaseJobTypeHandler: base,
	}
}

// ValidateJob validates a notebook job
func (h *NotebookJobTypeHandler) ValidateJob(job *mysql.Job) error {
	// Call base validation
	if err := h.BaseJobTypeHandler.ValidateJob(job); err != nil {
		return err
	}

	// Additional notebook-specific validation
	if job.Parameters == "" {
		return fmt.Errorf("notebook job requires parameters")
	}

	// TODO: Parse and validate notebook-specific parameters
	// This would include checking for required notebook_path parameter

	return nil
}

// CreateResources creates resources for a notebook job
func (h *NotebookJobTypeHandler) CreateResources(ctx context.Context, job *mysql.Job) ([]*job_manager.Resource, error) {
	// Create notebook-specific resources
	resources := []*job_manager.Resource{
		{
			Type:      "Pod",
			Name:      fmt.Sprintf("notebook-%s-pod", job.JobID),
			Namespace: "notebooks",
			Spec: map[string]interface{}{
				"containers": []map[string]interface{}{
					{
						"name":  "jupyter",
						"image": "jupyter/notebook:latest",
						"ports": []map[string]interface{}{
							{
								"containerPort": 8888,
								"protocol":      "TCP",
							},
						},
						"env": []map[string]interface{}{
							{
								"name":  "JUPYTER_ENABLE_LAB",
								"value": "yes",
							},
							{
								"name":  "JUPYTER_TOKEN",
								"value": "{{.Token}}",
							},
							{
								"name":  "JUPYTER_ALLOW_ORIGIN",
								"value": "{{.AllowOrigin}}",
							},
						},
						"resources": map[string]interface{}{
							"requests": map[string]interface{}{
								"cpu":    "{{.CPU}}",
								"memory": "{{.Memory}}",
							},
							"limits": map[string]interface{}{
								"cpu":    "{{.CPULimit}}",
								"memory": "{{.MemoryLimit}}",
							},
						},
						"volumeMounts": []map[string]interface{}{
							{
								"name":      "notebook-data",
								"mountPath": "/home/jovyan/work",
							},
						},
					},
				},
				"volumes": []map[string]interface{}{
					{
						"name": "notebook-data",
						"persistentVolumeClaim": map[string]interface{}{
							"claimName": fmt.Sprintf("notebook-%s-pvc", job.JobID),
						},
					},
				},
			},
			Status: map[string]interface{}{
				"phase": "Pending",
			},
			CreatedAt: time.Now(),
			UpdatedAt: time.Now(),
		},
		{
			Type:      "Service",
			Name:      fmt.Sprintf("notebook-%s-service", job.JobID),
			Namespace: "notebooks",
			Spec: map[string]interface{}{
				"ports": []map[string]interface{}{
					{
						"port":       8888,
						"targetPort": 8888,
						"protocol":   "TCP",
					},
				},
				"selector": map[string]interface{}{
					"app": fmt.Sprintf("notebook-%s", job.JobID),
				},
			},
			Status: map[string]interface{}{
				"type": "ClusterIP",
			},
			CreatedAt: time.Now(),
			UpdatedAt: time.Now(),
		},
		{
			Type:      "PersistentVolumeClaim",
			Name:      fmt.Sprintf("notebook-%s-pvc", job.JobID),
			Namespace: "notebooks",
			Spec: map[string]interface{}{
				"accessModes": []string{"ReadWriteOnce"},
				"resources": map[string]interface{}{
					"requests": map[string]interface{}{
						"storage": "10Gi",
					},
				},
			},
			Status: map[string]interface{}{
				"phase": "Pending",
			},
			CreatedAt: time.Now(),
			UpdatedAt: time.Now(),
		},
	}

	return resources, nil
}

// GetResourceTemplates gets resource templates for a notebook job
func (h *NotebookJobTypeHandler) GetResourceTemplates(job *mysql.Job) ([]*job_manager.ResourceTemplate, error) {
	templates := []*job_manager.ResourceTemplate{
		{
			Type:        "Pod",
			Name:        "notebook-pod-template",
			Namespace:   "notebooks",
			Description: "Jupyter notebook pod template",
			Template: map[string]interface{}{
				"apiVersion": "v1",
				"kind":       "Pod",
				"metadata": map[string]interface{}{
					"labels": map[string]interface{}{
						"app":      "jupyter-notebook",
						"job-id":   "{{.JobID}}",
						"job-type": "notebook",
					},
				},
				"spec": map[string]interface{}{
					"containers": []map[string]interface{}{
						{
							"name":  "jupyter",
							"image": "jupyter/notebook:latest",
							"ports": []map[string]interface{}{
								{
									"containerPort": "{{.Port}}",
									"protocol":      "TCP",
								},
							},
							"env": []map[string]interface{}{
								{
									"name":  "JUPYTER_ENABLE_LAB",
									"value": "yes",
								},
								{
									"name":  "JUPYTER_TOKEN",
									"value": "{{.Token}}",
								},
								{
									"name":  "JUPYTER_ALLOW_ORIGIN",
									"value": "{{.AllowOrigin}}",
								},
							},
							"resources": map[string]interface{}{
								"requests": map[string]interface{}{
									"cpu":    "{{.CPU}}",
									"memory": "{{.Memory}}",
								},
								"limits": map[string]interface{}{
									"cpu":    "{{.CPULimit}}",
									"memory": "{{.MemoryLimit}}",
								},
							},
						},
					},
				},
			},
			Parameters: h.supportedParameters,
		},
	}

	return templates, nil
}

// StartJob starts a notebook job
func (h *NotebookJobTypeHandler) StartJob(ctx context.Context, job *mysql.Job) error {
	// Notebook-specific startup logic
	// This could include:
	// - Setting up the notebook environment
	// - Loading initial notebooks
	// - Configuring authentication
	// - Setting up networking

	return nil
}

// StopJob stops a notebook job
func (h *NotebookJobTypeHandler) StopJob(ctx context.Context, job *mysql.Job) error {
	// Notebook-specific shutdown logic
	// This could include:
	// - Saving notebook state
	// - Graceful shutdown of Jupyter server
	// - Cleanup of temporary files

	return nil
}

// GetJobStatus gets the status of a notebook job
func (h *NotebookJobTypeHandler) GetJobStatus(ctx context.Context, job *mysql.Job) (*mysql.JobStatus, error) {
	// Notebook-specific status logic
	// This could include:
	// - Checking if Jupyter server is responding
	// - Getting notebook session information
	// - Checking resource usage

	status := &mysql.JobStatus{
		JobID:     job.JobID,
		Status:    string(job_manager.JobStateRunning),
		Message:   "Jupyter notebook is running",
		CreatedAt: time.Now(),
	}

	return status, nil
}
